{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import numpy.ma as ma\n",
    "\n",
    "import matplotlib.pyplot as plt \n",
    "import matplotlib\n",
    "from matplotlib import cm\n",
    "\n",
    "from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import\n",
    "from matplotlib.collections import PolyCollection\n",
    "from matplotlib import colors as mcolors\n",
    "\n",
    "from scipy.interpolate import CubicSpline, UnivariateSpline, griddata\n",
    "%matplotlib inline\n",
    "\n",
    "\n",
    "def bin_data(xi, yi):\n",
    "    x = np.unique(xi)\n",
    "    y = [None] * x.size\n",
    "    for i in range(x.size): y[i] = yi[xi == x[i]]\n",
    "    # Return \n",
    "    y_mean = np.array([a.mean() for a in y])\n",
    "    y_err = np.array([np.std(a) / (a.size**0.5) for a in y])\n",
    "    return (x, y_mean, y_err, y, xi, yi)\n",
    "\n",
    "\n",
    "\n",
    "def cc(arg):\n",
    "    '''\n",
    "    Shorthand to convert 'named' colors to rgba format at 60% opacity.\n",
    "    '''\n",
    "    return mcolors.to_rgba(arg, alpha=0.6)\n",
    "\n",
    "\n",
    "def polygon_under_graph(xlist, ylist):\n",
    "    '''\n",
    "    Construct the vertex list which defines the polygon filling the space under\n",
    "    the (xlist, ylist) line graph.  Assumes the xs are in ascending order.\n",
    "    '''\n",
    "    return [(xlist[0], 0.), *zip(xlist, ylist), (xlist[-1], 0.)]\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import colors as mcolors\n",
    "colorsraw = [\n",
    "    (10, 49, 244),\n",
    "#     (130,75,111),\n",
    "    (252, 30, 32)\n",
    "]\n",
    "colors = [[c/255 for c in cl] for cl in colorsraw]\n",
    "\n",
    "cmapbuild = mcolors.LinearSegmentedColormap.from_list('colormap', colors)\n",
    "\n",
    "def cmapscaled2(temp, cscale, alpha, saturation=0.7, lightness=0.9):\n",
    "#     return [*cmapbuild(temp/cscale)[:2],alpha]\n",
    "    temp=temp-.07\n",
    "    originalhsv = np.array(mcolors.rgb_to_hsv(cmapbuild(temp/cscale)[:3]))\n",
    "    modifier = np.array([0,saturation,lightness])\n",
    "    newrgb = mcolors.hsv_to_rgb(originalhsv-modifier)\n",
    "    return [*newrgb[:3], alpha]\n",
    "\n",
    "cmapscaled = cmapscaled2\n",
    "\n",
    "cscale=1\n",
    "\n",
    "fillcmap = lambda temp: cmapscaled(temp, cscale, 0.2, 0, 0.1)\n",
    "facecmap = lambda temp: cmapscaled(temp, cscale, 1, 0.4,.010)\n",
    "edgecmap = lambda temp: cmapscaled(temp, cscale, 1, 0.1,0.025)\n",
    "linecmap = lambda temp: cmapscaled(temp, cscale, 0.9,0.1,.025)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig,ax=plt.subplots(figsize=[2.25,3.3],nrows=1,ncols=1)\n",
    "\n",
    "# ax.set_position([0.12,0.14,0.76,0.85])\n",
    "\n",
    "# for j in range(len(index)):\n",
    "#     i=index[j]\n",
    "#     r=df_all.iloc[i]\n",
    "    \n",
    "    \n",
    "#     temperature='{:.2f}$T_c$'.format(r.TTilde/Tc)\n",
    "#     temp_text_string=temperature\n",
    "#     plt.text(150,-RF_space*j+0.33,temp_text_string,fontsize=8)\n",
    "#     offset_temp=0\n",
    "#     if df_all.Fit_method[i]=='Sinusoidal':\n",
    "#         plt.plot(xfit_show,fitfun_decay(xfit_show,*df_all.popt_sin[i][:-1])-RF_space*j,'--',color=linecmap(j/(len(index)-1)),alpha=0.7,linewidth=0.75)\n",
    "#         offset_temp=df_all.popt_sin[i][-1]\n",
    "#     elif df_all.Fit_method[i]=='Overdampen':\n",
    "#         plt.plot(xfit_show,fitfun_overdamp(xfit_show,*df_all.popt_overdamp[i][:-1])-RF_space*j,'--',color=linecmap(j/(len(index)-1)),alpha=0.7,linewidth=0.75)\n",
    "#         offset_temp=df_all.popt_overdamp[i][-1]\n",
    "#     else:\n",
    "#         plt.plot(xfit_show,fitfun_expdecay(xfit_show,*df_all.popt_exp[i][:-1])-RF_space*j,'--',color=linecmap(j/(len(index)-1)),alpha=0.7,linewidth=0.75)\n",
    "#         offset_temp=df_all.popt_exp[i][-1]\n",
    "        \n",
    "#     plt.plot(df_all.WaitTime[i],df_all.RF_evolutionN[i]-offset_temp-RF_space*j,'o',markersize=1.5,mew=0.5,mfc=facecmap(j/(len(index)-1)),mec=edgecmap(j/(len(index)-1)))\n",
    "    \n",
    "#     plt.axhline(-RF_space*j,linestyle=':',color=[0.3,0.3,0.3],alpha=0.4,linewidth=0.5)\n",
    "    \n",
    "#     temperature='T/Tc={:.3f}'.format(r.TTilde/Tc)\n",
    "#     Condensate='Cond={:.3f}'.format(r.Cond)\n",
    "#     temp_text_string=r.DS_name+', '+temperature+', '+Condensate\n",
    "#     plt.xlim([0,150])\n",
    "#     plt.grid()\n",
    "#     plt.ylabel('RF Response (a.u.)',fontsize=9,labelpad=0)\n",
    "# #     plt.xlabel('Time [ms]',fontsize=9)\n",
    "\n",
    "# ax.tick_params(axis='both',direction='in',length=2,top=1,right=1,labelsize=7)\n",
    "# ax.set_xlabel('t (ms)',fontsize=9,labelpad=1)\n",
    "# plt.yticks(np.linspace(-(len(index)-1)*RF_space,0,len(index)))\n",
    "# ax.set_yticklabels(['','','','','',''])\n",
    "\n",
    "# fig.savefig('LocalHeater_Waterfall_colorful.pdf',dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import cm\n",
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "\n",
    "bwr_cmap = cm.get_cmap('bwr', 256)\n",
    "\n",
    "a=1.0\n",
    "\n",
    "sclae=a*np.linspace(0, 1, 256)+(1-a)*(((np.linspace(0, 1, 256)-0.5)*2)**3/2+0.5)\n",
    "\n",
    "plt.plot(sclae)\n",
    "\n",
    "bwr_nc = bwr_cmap(sclae)\n",
    "bwr_nc = ListedColormap(bwr_nc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## All curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all=pd.read_pickle('FigS8.pkl')\n",
    "\n",
    "df_all['RF_evolutionNN']=[None]*len(df_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fitfun_decay(t,nu=20,A=1,B=0,gamma=1,C=0):\n",
    "    return np.exp(-gamma*2*np.pi*t/1000/2)*(A*np.cos(np.sqrt(nu**2-(gamma/2)**2)*2*np.pi*t/1000)+B*np.sin(np.sqrt(nu**2-(gamma/2)**2)*2*np.pi*t/1000))+C\n",
    "\n",
    "def fitfun_overdamp(t,nu=20,A=1,B=0,gamma=1,C=0):\n",
    "    return np.exp(-gamma*2*np.pi*t/1000/2)*(A*np.exp(-np.sqrt((gamma/2)**2-nu**2)*2*np.pi*t/1000)+B*np.exp(np.sqrt((gamma/2)**2-nu**2)*2*np.pi*t/1000))+C\n",
    "\n",
    "def fitfun_expdecay(t,alpha=20,A=1,offset=0):\n",
    "    return A*np.exp(-alpha*2*np.pi*t/1000)+offset\n",
    "\n",
    "index=[0,2,4,6,9,14]\n",
    "RF_space=-1.6\n",
    "Tc=0.167\n",
    "\n",
    "xfit_show=np.linspace(0,400,2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RF_space=1.75\n",
    "\n",
    "fig,axs=plt.subplots(ncols=3,nrows=1,figsize=[6,9],sharey=True)\n",
    "\n",
    "axs[0].tick_params(axis='both',direction='in',right=1,top=1,length=0,labelsize=9)\n",
    "axs[1].tick_params(axis='both',direction='in',right=1,top=1,length=0,labelsize=9)\n",
    "axs[2].tick_params(axis='both',direction='in',right=1,top=1,length=0,labelsize=9)\n",
    "\n",
    "\n",
    "axs[0].set_position([0.1,0.05,0.56,0.9])\n",
    "axs[1].set_position([0.66,0.05,0.16,0.9])\n",
    "axs[2].set_position([0.82,0.05,0.16,0.9])\n",
    "\n",
    "Dlen=len(df_all)\n",
    "\n",
    "axs[1].set_xticklabels([])\n",
    "axs[1].set_xlim([0,1])\n",
    "axs[2].set_xticklabels([])\n",
    "axs[2].set_xlim([0,1])\n",
    "plt.sca(axs[0])\n",
    "\n",
    "\n",
    "for i,r in (df_all.iterrows()):\n",
    "    plt.sca(axs[0])\n",
    "    para_text=\"\"\n",
    "    category_text=\"\"\n",
    "    if df_all.Fit_method[i]=='Sinusoidal':\n",
    "        offset=df_all.popt_sin[i][-1]\n",
    "        plt.plot(xfit_show,fitfun_decay(xfit_show,*df_all.popt_sin[i])+RF_space*i-offset,'--',color=linecmap(i/(Dlen-1)),linewidth=1)\n",
    "        para_text=para_text+r\"$c_2 k= {:.1f}({:.1f})$Hz,\".format(df_all.nu_fit[i],df_all.nu_fit_err[i])+\"\\n\"+r\"$D_2 k^2 = {:.1f}({:.1f})$Hz.\".format(df_all.Gamma_fit[i],df_all.Gamma_fit_err[i])\n",
    "        category_text=category_text+'Underdamped'\n",
    "    elif df_all.Fit_method[i]=='Overdampen':\n",
    "        offset=df_all.popt_overdamp[i][-1]\n",
    "        plt.plot(xfit_show,fitfun_overdamp(xfit_show,*df_all.popt_overdamp[i])+RF_space*i-offset,'--',color=linecmap(i/(Dlen-1)),linewidth=1)\n",
    "        para_text=para_text+r\"$c_2 k= {:.1f}({:.1f})$Hz,\".format(df_all.nu_fit[i],df_all.nu_fit_err[i])+\"\\n\"+r'$D_2 k^2 = {:.1f}({:.1f})$Hz.'.format(df_all.Gamma_fit[i],df_all.Gamma_fit_err[i])\n",
    "        category_text=category_text+'Overdamped'\n",
    "    else:\n",
    "        offset=df_all.popt_exp[i][-1]\n",
    "        plt.plot(xfit_show,fitfun_expdecay(xfit_show,*df_all.popt_exp[i])+RF_space*i-offset,'--',color=linecmap(i/(Dlen-1)),linewidth=1)\n",
    "        para_text=para_text+r'$D_2 k^2 = {:.1f}({:.1f})$Hz.'.format(df_all.Gamma_fit[i],df_all.Gamma_fit_err[i])\n",
    "        category_text=category_text+r'$c_2 k=0$'\n",
    "\n",
    "    plt.errorbar(df_all.WaitTime[i],df_all.RF_evolutionN[i]+RF_space*i-offset,yerr=df_all.rf_errN[i],marker='o',linestyle='None',markersize=2.5,\n",
    "                elinewidth=0.5,mew=0.2,mec=edgecmap(i/(Dlen-1)),mfc=facecmap(i/(Dlen-1)),ecolor=edgecmap(i/(Dlen-1)))\n",
    "    \n",
    "    \n",
    "    temperature='T/Tc={:.2f}({:.0f})'.format(r.TTilde/Tc,r.TTilde_err/Tc*100)\n",
    "    \n",
    "    temp_text_string=category_text+\"\\n\"+para_text\n",
    "#     plt.text(120,RF_space*i-RF_space*0.45,temperature,fontsize=9)\n",
    "    plt.xlim([0,200])\n",
    "    \n",
    "    \n",
    "    plt.ylabel(r'$\\Delta T$ (a.u.)',fontsize=9)\n",
    "    plt.xlabel(r'$t$ (ms)',fontsize=9)\n",
    "    \n",
    "    plt.sca(axs[1])\n",
    "    plt.text(0.035,RF_space*i-RF_space*0.1,temperature,fontsize=9,weight='bold')\n",
    "    \n",
    "    plt.sca(axs[2])\n",
    "    plt.text(0.035,RF_space*i-RF_space*0.3,temp_text_string,fontsize=7)\n",
    "\n",
    "plt.ylim([-RF_space,RF_space*len(df_all)])\n",
    "plt.yticks(np.arange(-RF_space,(RF_space*len(df_all))+RF_space,RF_space))\n",
    "axs[0].set_yticklabels([])\n",
    "plt.sca(axs[0])\n",
    "plt.grid(linewidth=0.25,alpha=0.5)\n",
    "\n",
    "\n",
    "\n",
    "fig.savefig('FigS9.pdf',dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RF_space=1.75\n",
    "\n",
    "fig,axs=plt.subplots(ncols=2,nrows=1,figsize=[6,9],sharey=True)\n",
    "\n",
    "axs[0].tick_params(axis='both',direction='in',right=1,top=1,length=0,labelsize=9)\n",
    "axs[1].tick_params(axis='both',direction='in',right=1,top=1,length=0,labelsize=9)\n",
    "\n",
    "axs[0].set_position([0.1,0.05,0.63,0.9])\n",
    "axs[1].set_position([0.73,0.05,0.25,0.9])\n",
    "\n",
    "axs[0].spines['right'].set_visible(False)\n",
    "axs[1].spines['left'].set_visible(False)\n",
    "\n",
    "Dlen=len(df_all)\n",
    "\n",
    "axs[1].set_xticklabels([])\n",
    "axs[1].set_xlim([0,1])\n",
    "plt.sca(axs[0])\n",
    "\n",
    "for i,r in (df_all.iterrows()):\n",
    "    text_position=0.45\n",
    "    plt.sca(axs[0])\n",
    "    para_text=\"\"\n",
    "    category_text=\"\"\n",
    "    if df_all.Fit_method[i]=='Sinusoidal':\n",
    "        offset=df_all.popt_sin[i][-1]\n",
    "        plt.plot(xfit_show,fitfun_decay(xfit_show,*df_all.popt_sin[i])+RF_space*i-offset,'--',color=linecmap(i/(Dlen-1)),linewidth=1)\n",
    "        para_text=r\"$c_2 k= {:.1f}({:.1f})$Hz,\".format(df_all.nu_fit[i],df_all.nu_fit_err[i])+\"\\n\"+r\"$D_2 k^2 = {:.1f}({:.1f})$Hz.\".format(df_all.Gamma_fit[i],df_all.Gamma_fit_err[i])\n",
    "        category_text=r'Underdamped,'\n",
    "    elif df_all.Fit_method[i]=='Overdampen':\n",
    "        offset=df_all.popt_overdamp[i][-1]\n",
    "        plt.plot(xfit_show,fitfun_overdamp(xfit_show,*df_all.popt_overdamp[i])+RF_space*i-offset,'--',color=linecmap(i/(Dlen-1)),linewidth=1)\n",
    "        para_text=r\"$c_2 k= {:.1f}({:.1f})$Hz,\".format(df_all.nu_fit[i],df_all.nu_fit_err[i])+\"\\n\"+r'$D_2 k^2 = {:.1f}({:.1f})$Hz.'.format(df_all.Gamma_fit[i],df_all.Gamma_fit_err[i])\n",
    "        category_text=r'Overdamped,'\n",
    "    else:\n",
    "        offset=df_all.popt_exp[i][-1]\n",
    "        plt.plot(xfit_show,fitfun_expdecay(xfit_show,*df_all.popt_exp[i])+RF_space*i-offset,'--',color=linecmap(i/(Dlen-1)),linewidth=1)\n",
    "        para_text=r'$D_2 k^2 = {:.1f}({:.1f})$Hz.'.format(df_all.Gamma_fit[i],df_all.Gamma_fit_err[i])\n",
    "        category_text=r'$c_2 k=0$,'\n",
    "        text_position=0.2\n",
    "\n",
    "    plt.errorbar(df_all.WaitTime[i],df_all.RF_evolutionN[i]+RF_space*i-offset,yerr=df_all.rf_errN[i],marker='o',linestyle='None',markersize=2.5,\n",
    "                elinewidth=0.5,mew=0.2,mec=edgecmap(i/(Dlen-1)),mfc=facecmap(i/(Dlen-1)),ecolor=edgecmap(i/(Dlen-1)))\n",
    "    \n",
    "    \n",
    "    temperature=r'T/Tc={:.2f}({:.0f}), '.format(r.TTilde/Tc,r.TTilde_err/Tc*100)\n",
    "    \n",
    "    temp_text_string=temperature+category_text+\"\\n\"+para_text\n",
    "#     plt.text(120,RF_space*i-RF_space*0.45,temperature,fontsize=9)\n",
    "    plt.xlim([0,200])\n",
    "    \n",
    "    \n",
    "    plt.ylabel(r'$\\Delta T$ (a.u.)',fontsize=9)\n",
    "    plt.xlabel(r'$t$ (ms)',fontsize=9)\n",
    "    \n",
    "    \n",
    "    plt.sca(axs[1])\n",
    "    plt.text(0.035,RF_space*i-RF_space*text_position,temp_text_string,fontsize=7)\n",
    "    \n",
    "plt.ylim([-RF_space,RF_space*len(df_all)])\n",
    "plt.yticks(np.arange(-RF_space,(RF_space*len(df_all))+RF_space,RF_space))\n",
    "axs[0].set_yticklabels([])\n",
    "plt.sca(axs[0])\n",
    "plt.grid(linewidth=0.25,alpha=0.5)\n",
    "\n",
    "\n",
    "\n",
    "fig.savefig('FigS9.pdf',dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i=16\n",
    "r=df_all.iloc[i]\n",
    "\n",
    "temperature=r'T/TF={:.3f}({:.0f}), '.format(r.TTilde,r.TTilde_err*1000)\n",
    "print(temperature)\n",
    "print(r.EF/1e3)\n",
    "plt.errorbar(df_all.WaitTime[i],df_all.RF_evolutionN[i],yerr=df_all.rf_errN[i],marker='o',linestyle='None',markersize=2.5,\n",
    "                elinewidth=0.5,mew=0.2,mec=edgecmap(i/(Dlen-1)),mfc=facecmap(i/(Dlen-1)),ecolor=edgecmap(i/(Dlen-1)))\n",
    "\n",
    "df_out=pd.DataFrame({'t (ms)':df_all.WaitTime[i],'Delta T':df_all.RF_evolutionN[i],'Delta T_err':df_all.rf_errN[i]})\n",
    "df_out.to_clipboard(index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {
    "height": "120px",
    "width": "252px"
   },
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "630px",
    "left": "0px",
    "right": "1112px",
    "top": "106px",
    "width": "212px"
   },
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
